import typing

import numpy as np
from PyQt5.QtCore import QRect

from yolo_detect_.YoloDetector import YoloDetector, RoiBox, Box
import cv2
from loguru import logger
from Tools.QtColorManager import QtColorManager
from Tools.DataStructs import DataStructsTools, ClassesLabel, RectInfo, PainterInfo


class YoloAutoMaker:
    def __init__(self, label_list, yolo_weight_path='/home/hello/下载/yolo11x-seg.pt'):
        self.label_list_self = label_list
        self.detector = YoloDetector(yolo_weight_path)

    def find_label_by_name(self, name) -> None | ClassesLabel:
        for _label in self.label_list_self:
            if _label.name == name:
                return _label
        logger.info('can not find label by id:{}'.format(str(name)))
        return None

    def get_mark_rect(self, _image: np.array) -> typing.List[RectInfo]:
        def to_qt_rect(roi_box: Box) -> QRect:
            return QRect(roi_box.x - roi_box.width // 2, roi_box.y - roi_box.height // 2, roi_box.width, roi_box.height)

        roi_box_list, _image, names = self.detector.detect(_image)
        rect_info_list: typing.List[RectInfo] = []
        for roi_box in roi_box_list:
            label = self.find_label_by_name(names[roi_box.class_id])
            if label is None:
                continue
            painter_info = PainterInfo(QtColorManager().get_color(label.index), label.name)
            # 保持显示的和用来处理的图像大小一致
            temp_rect: RectInfo = RectInfo(to_qt_rect(roi_box.box), label.index, painter_info)
            rect_info_list.append(temp_rect)
        return rect_info_list


if __name__ == '__main__':
    label_list = DataStructsTools.get_label_list_from_json(
        '/home/hello/graduated_ws/src/yolo_mark_tools/config/label.json')
    detector = YoloAutoMaker(label_list)
    image = cv2.imread('/home/hello/图片/yolo_data/images/train/截图 2024-06-27 10-53-07.png')
    print(detector.get_mark_rect(image))
